from torch.utils.data import  Dataset
from PIL import Image
import os
class MyData(Dataset):

    def __init__(self,root_dir,label_dir):
        #n
        self.root_dir = root_dir
        #ants
        self.label_dir = label_dir
        #train/ants
        self.path = os.path.join(self.root_dir,self.label_dir)
        # 显示path中文件的all list
        self.img_path = os.listdir(self.path)

    # 获取返回list的每一项内容
    def __getitem__(self, item):
        img_name = self.img_path[item]
        img_item_path = os.path.join(self.path,img_name)
        img = Image.open(img_item_path)
        label = self.label_dir
        return img , label

    # 返回list的长度
    def __len__(self):
        return len(self.img_path)

root_dir = "../data/edata/train"
ants_label = "ants"
bees_label = "bees"

ants_dataset = MyData(root_dir,ants_label)
bees_dataset = MyData(root_dir,bees_label)
train_dateset = ants_dataset + bees_dataset

img,label = train_dateset[200]
img.show()
